from __future__ import annotations

from typing import TYPE_CHECKING

import gymnasium as gym

if TYPE_CHECKING:
    from driving_gym.environment import DrivingGym


class SingleAgentWrapper(gym.Wrapper):
    def __init__(self, env: DrivingGym, target_actor: str = "hero"):
        if target_actor not in env.agents:
            raise ValueError(f"Target actor {target_actor} is not in the environment.")

        super(SingleAgentWrapper, self).__init__(env)
        self.actor = target_actor
        self.observation_space = env.observation_space[self.actor]
        self.action_space = env.action_space[self.actor]

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        return obs[self.actor], info

    def step(self, action):
        actions = self.env.action_space.sample()
        if isinstance(action, dict):
            actions.update(action)
        else:
            actions[self.actor] = action

        obs, reward, te, tr, info = self.env.step(actions)
        return obs[self.actor], reward[self.actor], te[self.actor], tr[self.actor], info
